#pragma once

#include "../Tradition/RHIPipeline.h"

enum RayTracingGroupType
{
	Raygen = 0,
	Miss = 1,
	Hit = 2,
	Callable = 3
};

class RayTracingLibrary : public RHIPipeline
{
	struct BuilderDetails {

		RefCountPtr<RHIDescriptor> descriptor;
	
		std::vector<ShaderPtr<RHIShader>> shaderStages;
		VkRayTracingShaderGroupCreateInfoKHR tempCroup;
		std::vector<VkRayTracingShaderGroupCreateInfoKHR> shaderGroups;
		uint8* dataPtr;
		uint64 dataSize;
		std::array<std::unordered_map<uint32, std::vector<uint8>>, 4> dataGroups;

		uint32 maxPipelineRayRecursionDepth;

		uint32 maxPipelineRayHitAttributeSize;

		uint32 maxPipelineRayPayloadSize;

		BuilderDetails()
		{
			tempCroup = {};
			tempCroup.sType = VK_STRUCTURE_TYPE_RAY_TRACING_SHADER_GROUP_CREATE_INFO_KHR;
			tempCroup.type = VK_RAY_TRACING_SHADER_GROUP_TYPE_GENERAL_KHR;
			tempCroup.generalShader = VK_SHADER_UNUSED_KHR;
			tempCroup.closestHitShader = VK_SHADER_UNUSED_KHR;
			tempCroup.anyHitShader = VK_SHADER_UNUSED_KHR;
			tempCroup.intersectionShader = VK_SHADER_UNUSED_KHR;

			dataPtr = nullptr;
			dataSize = 0;

			maxPipelineRayRecursionDepth = 0;

			maxPipelineRayHitAttributeSize = 0;
			maxPipelineRayPayloadSize = 0;
		}
	};

public:

	class Builder : public BuilderBase<BuilderDetails> {
		friend struct BuilderDetails;
	public:
		Builder() noexcept;
		Builder(Builder const& rhs) noexcept;
		Builder(Builder&& rhs) noexcept;
		~Builder() noexcept;
		Builder& operator=(Builder const& rhs) noexcept;
		Builder& operator=(Builder&& rhs) noexcept;

		/************************************************************************/
		/*                              Descriptor                              */
		/************************************************************************/

		Builder& SetDescriptor(RefCountPtr<RHIDescriptor> descriptor) noexcept;

		/************************************************************************/
		/*                           ShaderStage                                */
		/************************************************************************/

		Builder& AddShaderStage(ShaderPtr<RHIShader> shader) noexcept;

		/************************************************************************/
		/*                           ShaderGroup                                */
		/************************************************************************/

		Builder& PushShader(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushRaygen(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushMiss(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushClosestHit(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushAnyHit(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushIntersection(ShaderPtr<RHIShader> shader) noexcept;
		Builder& PushCallable(ShaderPtr<RHIShader> shader) noexcept;
		template<typename T>
		Builder& PushData(T& data) noexcept
		{
			PushData((uint8*)&data, sizeof(T));
			return *this;
		}
		Builder& PackageGroup() noexcept;

		/************************************************************************/
		/*                           Ray   Param                                */
		/************************************************************************/

		Builder& SetMaxPipelineRayRecursionDepth(uint32 maxPipelineRayRecursionDepth) noexcept;

		/************************************************************************/
		/*                          PipelineInterface                           */
		/************************************************************************/

		Builder& SetMaxPipelineRayHitAttributeSize(uint32 maxPipelineRayHitAttributeSize) noexcept;
		Builder& SetMaxPipelineRayPayloadSize(uint32 maxPipelineRayPayloadSize) noexcept;

		RefCountPtr<RayTracingLibrary> Build();
		RefCountPtr<RayTracingLibrary> Build(RenderProcess* process, const std::string& name);

	protected:

		uint32 GetShaderIndex(ShaderPtr<RHIShader> shader) const;

		void PushData(uint8* data, uint64 dataSize);

	private:

		friend class RayTracingLibrary;
	};

	RayTracingLibrary(BuilderDetails& builderDetails);

	virtual void ReleaseResource() override;

	std::vector<VkRayTracingShaderGroupCreateInfoKHR>& GetShaderGroups();

	std::array<std::unordered_map<uint32, std::vector<uint8>>, 4>& GetDataGroups();

	VkRayTracingPipelineInterfaceCreateInfoKHR GetRayTracingPipelineInterface() const;

	uint32 GetMaxPipelineRayRecursionDepth() const;

	std::vector<ShaderPtr<RHIShader>>& GetShaderStages();

protected:

	std::vector<VkRayTracingShaderGroupCreateInfoKHR> mShaderGroups;

	std::array<std::unordered_map<uint32, std::vector<uint8>>, 4> mDataGroups;

	VkRayTracingPipelineInterfaceCreateInfoKHR mRayTracingPipelineInterface;

	uint32 mMaxPipelineRayRecursionDepth;

	std::vector<ShaderPtr<RHIShader>> mShaderStages;
};